package com.qezhhnjy.leetcode.question.tree;

import com.qezhhnjy.leetcode.entity.TreeNode;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author qezhhnjy
 * @date 2020/12/20-17:55
 * 给定一个非空二叉树, 返回一个由每层节点平均值组成的数组。
 * <p>
 * 深度优先搜索：Depth First Search
 * 广度优先搜索：Breadth First Search
 */
public class BinaryTreeLevelAverageValue {

    public List<Double> list = new ArrayList<>();
    public List<TreeNode> level = new ArrayList<>();

    public List<Double> averageOfLevels(TreeNode root) {
        if (root == null) return null;
        level.add(root);

        while (!level.isEmpty()) {
            list.add(level.stream().mapToInt(node -> node.val).average().getAsDouble());

            level = level.stream().flatMap(node -> Stream.of(node.left, node.right))
                    .filter(Objects::nonNull).collect(Collectors.toList());
        }
        return list;
    }

    public LinkedList<TreeNode> queue = new LinkedList<>();

    public List<Double> breadth(TreeNode root) {
        if (root == null) return null;
        queue.add(root);
        while (!queue.isEmpty()) {
            double sum = 0;
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                TreeNode poll = queue.poll();
                sum += poll.val;

                if (poll.left != null) queue.add(poll.left);
                if (poll.right != null) queue.add(poll.right);
            }
            list.add(sum / size);
        }
        return list;
    }

    public List<Double> depth(TreeNode root) {
        List<Integer> counts = new ArrayList<>();
        List<Double> sums = new ArrayList<>();
        dfs(root, 0, counts, sums);
        List<Double> averages = new ArrayList<>();
        int size = sums.size();
        for (int i = 0; i < size; i++) {
            averages.add(sums.get(i) / counts.get(i));
        }
        return averages;
    }

    public void dfs(TreeNode root, int level, List<Integer> counts, List<Double> sums) {
        if (root == null) {
            return;
        }
        if (level < sums.size()) {
            sums.set(level, sums.get(level) + root.val);
            counts.set(level, counts.get(level) + 1);
        } else {
            sums.add(1.0 * root.val);
            counts.add(1);
        }
        dfs(root.left, level + 1, counts, sums);
        dfs(root.right, level + 1, counts, sums);
    }
}
